﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using AIOptimization.Algorithms.Repositories;

namespace AIOptimization.Algorithms
{
    public class AntColonyOptimization
    {
        Random rd = new Random(0);
        bool parameterSeted = false;
        int Alpha;
        int Beta;

        double Rho;
        double Q;

        int NumCities;
        int NumAnts;
        int MaxTime;
        public void SetParameter(int numCities, int numAnts, int maxTime, double rho = 0.01, double q = 2.0, int alpha = 3, int beta = 2)
        {
            parameterSeted = true;
            NumCities = numCities;
            NumAnts = numAnts;
            MaxTime = maxTime;
            Alpha = alpha;
            Beta = beta;
            Rho = rho;
            Q = q;
        }

        public Ant Execute()
        {
            if (!parameterSeted) throw new Exception("You must be set parameter. Using SetParameter Method");
            Ant result = new Ant();
            int[][] dists = MakeGraphDistances(NumCities);
            int[][] ants = InitializeAnt(NumAnts, NumCities);
            result.BestTrail = BestTrail(ants, dists);
            result.BestLength = Length(result.BestTrail, dists);

            double[][] pheromones = InitializePheromones(NumCities);

            int time = 0;
            result.TrainResult.Add(new Train(time, result.BestLength, result.BestTrail));

            while (time < MaxTime)
            {
                UpdateAnts(ants, pheromones, dists);
                UpdatePheromones(pheromones, ants, dists);

                result.TrainResult.Add(new Train(time, Length(BestTrail(ants, dists), dists), BestTrail(ants, dists)));
                if (result.TrainResult.Last().Length < result.BestLength)
                {
                    result.BestLength = result.TrainResult.Last().Length;
                    result.BestTrail = result.TrainResult.Last().Trail;
                    result.BestTime = time;
                }
                ++time;
            }
            return result;
        }

        double[][] InitializePheromones(int numCities)
        {
            double[][] pheromones = new double[numCities][];
            for (int i = 0; i < numCities; ++i)
                pheromones[i] = new double[numCities];
            for (int i = 0; i < pheromones.Length; ++i)
                for (int j = 0; j < pheromones[i].Length; ++j)
                    pheromones[i][j] = 0.01; // otherwise first call to UpdateAnts -> BuiuldTrail -> NextNode -> MoveProbs => all 0.0 => throws
            return pheromones;
        }

        int[][] InitializeAnt(int numAnts, int numCities)
        {
            int[][] ants = new int[numAnts][];
            for (int k = 0; k < numAnts; ++k)
            {
                int start = rd.Next(0, numCities);
                ants[k] = RandomTrail(start, numCities);
            }
            return ants;
        }

        int[] RandomTrail(int start, int numCities)
        {
            int[] trail = new int[numCities];

            for (int i = 0; i < numCities; ++i) { trail[i] = i; }

            for (int i = 0; i < numCities; ++i)
            {
                int r = rd.Next(i, numCities);
                int tmp = trail[r]; trail[r] = trail[i]; trail[i] = tmp;
            }

            int idx = IndexOfTarget(trail, start);
            int temp = trail[0];
            trail[0] = trail[idx];
            trail[idx] = temp;

            return trail;
        }

        int IndexOfTarget(int[] trail, int target)
        {
            for (int i = 0; i < trail.Length; ++i)
            {
                if (trail[i] == target)
                    return i;
            }
            throw new Exception("Target not found in IndexOfTarget");
        }

        private int[][] MakeGraphDistances(int numCities)
        {
            int[][] dists = new int[numCities][];
            for (int i = 0; i < dists.Length; ++i)
                dists[i] = new int[numCities];
            for (int i = 0; i < numCities; ++i)
                for (int j = i + 1; j < numCities; ++j)
                {
                    int d = rd.Next(1, 9); // [1,8]
                    dists[i][j] = d;
                    dists[j][i] = d;
                }
            return dists;
        }

        int[] BestTrail(int[][] ants, int[][] dists)
        {
            double bestLength = Length(ants[0], dists);
            int idxBestLength = 0;
            for (int k = 1; k < ants.Length; ++k)
            {
                double len = Length(ants[k], dists);
                if (len < bestLength)
                {
                    bestLength = len;
                    idxBestLength = k;
                }
            }
            int numCities = ants[0].Length;
            int[] bestTrail = new int[numCities];
            ants[idxBestLength].CopyTo(bestTrail, 0);
            return bestTrail;
        }

        double Length(int[] trail, int[][] dists) // total length of a trail
        {
            double result = 0.0;
            for (int i = 0; i < trail.Length - 1; ++i)
                result += Distance(trail[i], trail[i + 1], dists);
            return result;
        }

        double Distance(int cityX, int cityY, int[][] dists)
        {
            return dists[cityX][cityY];
        }

        void UpdateAnts(int[][] ants, double[][] pheromones, int[][] dists)
        {
            int numCities = pheromones.Length;
            for (int k = 0; k < ants.Length; ++k)
            {
                int start = rd.Next(0, numCities);
                int[] newTrail = BuildTrail(k, start, pheromones, dists);
                ants[k] = newTrail;
            }
        }

        int[] BuildTrail(int k, int start, double[][] pheromones, int[][] dists)
        {
            int numCities = pheromones.Length;
            int[] trail = new int[numCities];
            bool[] visited = new bool[numCities];
            trail[0] = start;
            visited[start] = true;
            for (int i = 0; i < numCities - 1; ++i)
            {
                int cityX = trail[i];
                int next = NextCity(k, cityX, visited, pheromones, dists);
                trail[i + 1] = next;
                visited[next] = true;
            }
            return trail;
        }

        int NextCity(int k, int cityX, bool[] visited, double[][] pheromones, int[][] dists)
        {
            double[] probs = MoveProbs(k, cityX, visited, pheromones, dists);

            double[] cumul = new double[probs.Length + 1];
            for (int i = 0; i < probs.Length; ++i)
                cumul[i + 1] = cumul[i] + probs[i];

            double p = rd.NextDouble();

            for (int i = 0; i < cumul.Length - 1; ++i)
                if (p >= cumul[i] && p < cumul[i + 1])
                    return i;
            throw new Exception("Failure to return valid city in NextCity");
        }

        double[] MoveProbs(int k, int cityX, bool[] visited, double[][] pheromones, int[][] dists)
        {
            int numCities = pheromones.Length;
            double[] taueta = new double[numCities];
            double sum = 0.0;
            for (int i = 0; i < taueta.Length; ++i)
            {
                if (i == cityX)
                    taueta[i] = 0.0;
                else if (visited[i] == true)
                    taueta[i] = 0.0;
                else
                {
                    taueta[i] = Math.Pow(pheromones[cityX][i], Alpha) * Math.Pow((1.0 / Distance(cityX, i, dists)), Beta);
                    if (taueta[i] < 0.0001)
                        taueta[i] = 0.0001;
                    else if (taueta[i] > (double.MaxValue / (numCities * 100)))
                        taueta[i] = double.MaxValue / (numCities * 100);
                }
                sum += taueta[i];
            }

            double[] probs = new double[numCities];
            for (int i = 0; i < probs.Length; ++i)
                probs[i] = taueta[i] / sum;
            return probs;
        }

        void UpdatePheromones(double[][] pheromones, int[][] ants, int[][] dists)
        {
            for (int i = 0; i < pheromones.Length; ++i)
            {
                for (int j = i + 1; j < pheromones[i].Length; ++j)
                {
                    for (int k = 0; k < ants.Length; ++k)
                    {
                        double length = Length(ants[k], dists);
                        double decrease = (1.0 - Rho) * pheromones[i][j];
                        double increase = 0.0;
                        if (EdgeInTrail(i, j, ants[k]) == true) increase = (Q / length);

                        pheromones[i][j] = decrease + increase;

                        if (pheromones[i][j] < 0.0001)
                            pheromones[i][j] = 0.0001;
                        else if (pheromones[i][j] > 100000.0)
                            pheromones[i][j] = 100000.0;

                        pheromones[j][i] = pheromones[i][j];
                    }
                }
            }
        }

        bool EdgeInTrail(int cityX, int cityY, int[] trail)
        {
            int lastIndex = trail.Length - 1;
            int idx = IndexOfTarget(trail, cityX);

            if (idx == 0 && trail[1] == cityY) return true;
            else if (idx == 0 && trail[lastIndex] == cityY) return true;
            else if (idx == 0) return false;
            else if (idx == lastIndex && trail[lastIndex - 1] == cityY) return true;
            else if (idx == lastIndex && trail[0] == cityY) return true;
            else if (idx == lastIndex) return false;
            else if (trail[idx - 1] == cityY) return true;
            else if (trail[idx + 1] == cityY) return true;
            else return false;
        }
    }
}
